#!/usr/bin/env python


#Created on Sun Sep 13 15:34:55 2015

#author: ZhongCheng


import math
import numpy as np
import argparse
import os
import re
class periodStruc():
    def __init__(self,inputFile,inputType):
        self.inputType=inputType
        self.outputType=""
        self.inputFile=inputFile
        self.elementDict={1:"H",3:"Li",4:"Be",5:"B",6:"C",7:"N",8:"O",9:"F",14:"Si",15:"P",16:"S",17:"Cl",35:"Br",35:"I"}
        self.fileName=''.join(inputFile.split('.')[:-1]) or inputFile
        self.outputFile=self.fileName+'.cif'        
        self.fract_coord = []
        self.coord = []
        if inputType == "":
            inputType=inputFile.split('.')[-1]
        if inputType == 'gui':
            self.readgui()
        if inputType == 'STRUC':
            self.readStruc()
        if inputType == 'CONTCAR':
            self.readcontcar()
        if os.path.basename(inputFile) == 'CONTCAR' or os.path.basename(inputFile) == 'POSCAR':
            self.readcontcar()
    #read gui file from crystal14
    def readgui(self):
        guiFile=open(self.inputFile)
        lines = guiFile.readlines()
        self.a_axis_cart=map(float,lines[1].split())
        self.b_axis_cart=map(float,lines[2].split())
        self.c_axis_cart=map(float,lines[3].split())
        self.coord=[map(float,i.split()) for i in  lines[10:]]
        for line in self.coord:
            line[0]=self.elementDict[line[0]]
    #read STRUC file from crystal14
    def readStruc(self):
        strucFile = open(self.inputFile)
        lines = strucFile.readlines()
        self.a_axis_cart=[float(i)*0.529177249 for i in lines[1].split()]
        self.b_axis_cart=[float(i)*0.529177249 for i in lines[2].split()]
        self.c_axis_cart=[float(i)*0.529177249 for i in lines[3].split()]
    #    print lines[5:]
        self.coord=[i.split()[:1]+[float(i)*0.529177249 for i in i.split()[1:4]] for i in  lines[5:] if len(i.split()) == 5]
    #convert cell axis from cartesian coordinates to cell parameter
    def readcontcar(self):
        strucFile = open(self.inputFile)
        lines = strucFile.readlines()
        scaling_factor=float(lines[1].strip())
        self.a_axis_cart = np.array(map(float,lines[2].split()))*scaling_factor
        self.b_axis_cart = np.array(map(float,lines[3].split()))*scaling_factor
        self.c_axis_cart = np.array(map(float,lines[4].split()))*scaling_factor
        element = lines[5].strip().split()
        num_of_element = lines[6].strip().split()
        self.fract_coord = []
        if lines[7].strip().lower().startswith('c'):
            for i in range(len(element)):
                self.coord += [[element[i]] for n in range(int(num_of_element[i]))]
            i=0
            for n in lines[7:]:
                if re.search('-?\d+.\d+\s+-?\d+.\d+\s+-?\d+.\d+',n):
                    self.coord[i] += np.array(map(float,n.strip().split()))*scaling_factor 
                    i = i+1
        else:
            for i in range(len(element)):
                self.fract_coord += [[element[i]] for n in range(int(num_of_element[i]))]
            i=0
            for n in lines[7:]:
                if re.search('-?\d+.\d+\s+-?\d+.\d+\s+-?\d+.\d+',n):
                    self.fract_coord[i] += self.fract_coord[i]
                    self.fract_coord[i][0] = self.fract_coord[i][0]+str(i)
                    self.fract_coord[i] += map(float,n.strip().split()) 
                    i = i+1
    def cart2cell(self):
        self.cell_length_a=sum([i**2 for i in self.a_axis_cart])**0.5
        self.cell_length_b=sum([i**2 for i in self.b_axis_cart])**0.5
        self.cell_length_c=sum([i**2 for i in self.c_axis_cart])**0.5
        self.cell_angle_alpha=math.degrees(math.acos(sum([i*j for (i,j) in zip(self.b_axis_cart,self.c_axis_cart)])/(self.cell_length_b*self.cell_length_c)))
        self.cell_angle_beta=math.degrees(math.acos(sum([i*j for (i,j) in zip(self.a_axis_cart,self.c_axis_cart)])/(self.cell_length_a*self.cell_length_c)))   
        self.cell_angle_gamma=math.degrees(math.acos(sum([i*j for (i,j) in zip(self.a_axis_cart,self.b_axis_cart)])/(self.cell_length_a*self.cell_length_b)))
    #convert atom coordinates from cartesian to fractional, the format of cartesian coordinates is a 2D list 
    #e.g. [['C',2.3,2.2,2.5],['H',3.5,4.5,1.3]]
    def cart2fract(self):
        if self.fract_coord:
            return
        self.fract_coord=[]
        #construct rotation matrix to align a to x and b to xy plane
        #alpha=-math.atan(self.a_axis_cart[1]/self.a_axis_cart[0])
        #beta=math.atan(self.a_axis_cart[2]/(self.a_axis_cart[0]**2+self.a_axis_cart[1]**2)**0.5)
        ##rotation along z axis
        #Rz=np.matrix([[math.cos(alpha),-math.sin(alpha),0],[math.sin(alpha),math.cos(alpha),0],[0,0,1]])
        ##rotation along y axis
        #Ry=np.matrix([[math.cos(beta),0,math.sin(beta)],[0,1,0],[-math.sin(beta),0,math.cos(beta)]])
        ##rotation along x axis to make b align to xy plane
        #new_b_axis=np.array(np.dot(Ry*Rz,np.array(self.b_axis_cart))).reshape(-1,)
        #gamma=-math.atan(new_b_axis[2]/new_b_axis[1])
        #Rx=np.matrix([[1,0,0],[0,math.cos(gamma),-math.sin(gamma)],[0,math.sin(gamma),math.cos(gamma)]])
        #rotat_m=Rx*Ry*Rz
        ##construct conversion matrix from cartesian to fraction coordinates
        #a=self.cell_length_a
        #b=self.cell_length_b
        #c=self.cell_length_c
        #alpha=math.radians(self.cell_angle_alpha)
        #beta=math.radians(self.cell_angle_beta)
        #gamma=math.radians(self.cell_angle_gamma)
        ##v is the volume of a unit parallelepiped         
        #v=(1-math.cos(alpha)**2-math.cos(beta)**2-math.cos(gamma)**2+2*math.cos(alpha)*math.cos(beta)*math.cos(gamma))**0.5
        #row1=[1/a,-math.cos(gamma)/math.sin(gamma)/a,(math.cos(alpha)*math.cos(gamma)-math.cos(beta))/math.sin(gamma)/a/v]
        #row2=[0,1/math.sin(gamma)/b,(math.cos(beta)*math.cos(gamma)-math.cos(alpha))/math.sin(gamma)/b/v]
        #row3=[0,0,math.sin(gamma)/c/v]
        #conver_m=np.matrix([row1,row2,row3])
        #self.rot_conv_m=conver_m*rotat_m
        ##print self.rot_conv_m        
        #element_num=1
        #for coord in self.coord:
        #    fract_coord=np.dot(self.rot_conv_m,np.array(coord[1:])).tolist()[0]
        #    self.fract_coord.append([coord[0]+str(element_num),coord[0]]+fract_coord)
        #    element_num=element_num+1
        #print self.fract_coord
        #my old code:
        element_num=1
        for coord in self.coord:
            fract_coord=np.array(np.matrix(coord[1:])*np.matrix([self.a_axis_cart,self.b_axis_cart,self.c_axis_cart])**-1).tolist()[0]
            self.fract_coord.append([coord[0]+str(element_num),coord[0]]+fract_coord)
            element_num=element_num+1
            
    
    def writecif(self):
        cifFile=open(self.outputFile,'w')
        cifFile.write('data_'+self.fileName+'\n')
        cifFile.write('{:35s}{:.4f}\n'.format('_cell_length_a',self.cell_length_a)+
                      '{:35s}{:.4f}\n'.format('_cell_length_b',self.cell_length_b)+
                      '{:35s}{:.4f}\n'.format('_cell_length_c',self.cell_length_c)+
                      '{:35s}{:.4f}\n'.format('_cell_angle_alpha',self.cell_angle_alpha)+
                      '{:35s}{:.4f}\n'.format('_cell_angle_beta',self.cell_angle_beta)+
                      '{:35s}{:.4f}\n'.format('_cell_angle_gamma',self.cell_angle_gamma))
        cifFile.write('loop_\n'
                      '_atom_site_label\n'
                      '_atom_site_type_symbol\n'
                      '_atom_site_fract_x\n'
                      '_atom_site_fract_y\n'
                      '_atom_site_fract_z\n')
        for coord in self.fract_coord:
            cifFile.write('{:7s}{:7s}{:10.5f}{:10.5f}{:10.5f}\n'.format(*coord))
        cifFile.close()
        print self.outputFile+" generated"              
 
parser=argparse.ArgumentParser(formatter_class=argparse.RawTextHelpFormatter,description="A script used to generate cif file from other file type (now support .gui and .STRUC file from CRYSTAL14.")
parser.add_argument("-t","--type",help="designate a file type (gui,STRUC,CONTCAR), default: judge by the suffix of the input file",  type=str,default='')
parser.add_argument("inputStruc",nargs='+')
args=parser.parse_args()
for inputStruc in args.inputStruc:      
    struc=periodStruc(inputStruc,args.type)
    #struc.readgui()
    struc.cart2cell()
    struc.cart2fract()
    struc.writecif()
